# NiN块是NiN中的基础块。
# 它由一个卷积层加两个充当全连接层的1×1卷积层串联而成。
# 其中第一个卷积层的超参数可以自行设置，而第二和第三个卷积层的超参数一般是固定的。
import time
import torch
from torch import nn, optim
from basic_knowledge import d2lzh_pytorch as d2l
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')


def nin_block(in_channels, out_channels, kernel_size, stride, padding):
    blk = nn.Sequential(nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding),
                        nn.ReLU(),
                        nn.Conv2d(out_channels, out_channels, kernel_size=1),
                        nn.ReLU(),
                        nn.Conv2d(out_channels, out_channels, kernel_size=1),
                        nn.ReLU())
    return blk


# NiN是在AlexNet问世不久后提出的。它们的卷积层设定有类似之处。
# NiN使用卷积窗口形状分别为11×11、5×5和3×3的卷积层，
# 相应的输出通道数也与AlexNet中的一致。每个NiN块后接一个步幅为2、窗口形状为3×3的最大池化层。
# 除使用NiN块以外，NiN还有一个设计与AlexNet显著不同：NiN去掉了AlexNet最后的3个全连接层，
# 取而代之地，NiN使用了输出通道数等于标签类别数的NiN块，
# 然后使用全局平均池化层对每个通道中所有元素求平均并直接用于分类。
# 这里的形状全局平均池化层即窗口等于输入空间维形状的平均池化层。
# NiN的这个设计的好处是可以显著减小模型参数尺寸，
# 从而缓解过拟合。然而，该设计有时会造成获得有效模型的训练时间的增加。
net = nn.Sequential(
    nin_block(1, 96, kernel_size=11, stride=4, padding=0),
    nn.MaxPool2d(kernel_size=3, stride=2),
    nin_block(96, 256, kernel_size=5, stride=1, padding=2),
    nn.MaxPool2d(kernel_size=3, stride=2),
    nin_block(256, 384, kernel_size=3, stride=1, padding=1),
    nn.MaxPool2d(kernel_size=3, stride=2),
    nn.Dropout(0.5),
    # 标签类别数是10
    nin_block(384, 10, kernel_size=3, stride=1, padding=1),
    d2l.GlobalAvgPool2d(),
    # 将四维的输出转成二维的输出，其形状为(批量大小, 10)
    d2l.FlattenLayer())


batch_size = 64
# 如出现“out of memory”的报错信息，可减小batch_size或resize
train_iter, test_iter = d2l.load_data_fashion_mnist(batch_size, resize=224)

lr, num_epochs = 0.002, 5
optimizer = torch.optim.Adam(net.parameters(), lr=lr)
d2l.train_ch5(net, train_iter, test_iter, batch_size, optimizer, device, num_epochs)